package udpserver

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"github.com/kataras/golog"
	"math/rand"
	"net"
	"nprobe/core"
	"nprobe/enviroments"
	"nprobe/utils/routines"
	"time"
)

type UDPServerReceiver struct {
	core.BaseReceiver
	SrcIPSet map[string]bool
}

func NewUDPServerReceiver(
	ctx context.Context, name string, id int, task *core.ProbeTask, salts map[int][]byte, targetSet map[string]bool,
) *UDPServerReceiver {
	ppRBuf := core.NewPPRBuf(task.SpanMs, task.TimeoutMs)

	srcIPSet := make(map[string]bool)
	for _, srcIP := range task.SrcIPs {
		srcIPSet[srcIP] = true
	}

	return &UDPServerReceiver{
		BaseReceiver: core.BaseReceiver{
			Ctx:      ctx,
			Name:     name,
			ID:       id,
			Task:     task,
			RBuf:     ppRBuf,
			Salts:    salts,
			DstIPSet: targetSet,
		},
		SrcIPSet: srcIPSet,
	}
}

func (r *UDPServerReceiver) StartReceiver() error {
	routines.GoWithRecover(r.Name, r.doReceiveUDP)
	return nil
}

func (r *UDPServerReceiver) doReceiveUDP() {
	conn, err := core.OpenReceiverIP4PacketConn("ip4:udp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	filter := core.CreateBPFUDPServer(r.Task.SrcPortMin, r.Task.SrcPortMax, r.Task.DstPortMin, r.Task.DstPortMax)
	err = core.SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	r.CurrSpanRecord = core.NewSpanRecord(0)
	payloadBuf := core.AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			r.StopReceiver()
			return
		default:
			r.receiveOneUDPPkt(payloadBuf)
		}

	}
}

func (r *UDPServerReceiver) receiveOneUDPPkt(payloadBuf []byte) {

	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		return
	}

	rxTs := time.Now().UnixNano()
	if h.Version != 4 {
		return
	}

	remoteIP := h.Src.To4().String()
	localIP := h.Dst.To4().String()
	tos := uint8(h.TOS)
	if len(r.DstIPSet) > 0 && !r.DstIPSet[localIP] {
		return
	}

	if len(r.SrcIPSet) > 0 && !r.SrcIPSet[remoteIP] {
		return
	}

	r.processUDPEchoReply(localIP, remoteIP, tos, rxTs, payload)
}

func (r *UDPServerReceiver) processUDPEchoReply(localIP string, remoteIP string, tos uint8, rxTs int64, payload []byte) {
	// udp header
	srcPort := binary.BigEndian.Uint16(payload[0:2])
	dstPort := binary.BigEndian.Uint16(payload[2:4])

	// udp payload
	udpLoad := payload[8:]
	rxTxTs := int64(binary.BigEndian.Uint64(udpLoad[12:20]))
	seq := binary.BigEndian.Uint16(udpLoad[20:22])

	ipID := uint16(rand.Intn(1 << 16))
	msg, err := core.EncodeIPUDPPacket(localIP, remoteIP, ipID, tos, r.Task.TTL, dstPort, srcPort, udpLoad)
	if err != nil {
		return
	}

	err = r.Conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	if err != nil {
		return
	}
	_, err = r.Conn.WriteToIP(msg, &net.IPAddr{IP: net.ParseIP(remoteIP)})
	if err != nil {
		return
	}

	txTs := time.Now().UnixNano()

	bitFlip := false
	if r.Task.BitFlip && bytes.Equal(udpLoad[:8], []byte(enviroments.MagicNumber)) {
		bitFlip = !bytes.Equal(r.Salts[int(seq)%4][22:], udpLoad[22:])
		if bitFlip {
			golog.Debugf("salt: %x", r.Salts[int(seq)%4])
			golog.Debugf("data: %x", payload)
			golog.Debugf("[bitflip] src=%s  dst=%s seq=%d, ts=%d", localIP, remoteIP, seq, rxTxTs)
		}
	}

	r.recordPkt(localIP, remoteIP, srcPort, dstPort, seq, txTs, rxTs, bitFlip)
}

func (r *UDPServerReceiver) recordPkt(localIP string, remoteIP string, srcPort uint16, dstPort uint16, seq uint16, txTs int64, rxTs int64, bitFlip bool) {
	spanID := txTs / (r.Task.SpanMs * int64(time.Millisecond))
	golog.Debugf("%s: %s:%d->%s:%d, seq: %d, spanID: %d, rxTs: %d, txTs: %d, delta: %d, bitFlip: %t",
		r.Name, remoteIP, srcPort, localIP, dstPort, seq, spanID, rxTs, txTs, txTs-rxTs, bitFlip)

	if spanID > r.CurrSpanRecord.SpanID {
		r.CurrSpanRecord = core.NewSpanRecord(spanID)
		r.RBuf.Enqueue(r.CurrSpanRecord)
	}

	key := fmt.Sprintf("%s-%d", remoteIP, seq)
	r.CurrSpanRecord.PktRecordMap[key] = &core.PktRecord{
		Key:      key,
		SrcIP:    remoteIP,
		DstIP:    localIP,
		SrcPort:  srcPort,
		DstPort:  dstPort,
		Proto:    r.Task.Proto,
		Tos:      r.Task.TOS,
		RxTxTs:   txTs,
		TxTs:     rxTs,
		RxTs:     txTs,
		Received: true,
		BitFlip:  bitFlip,
	}
}
